import numpy as np
import pandas as pd
import math

#计算对应库
def cal_radar_gauge(d, beta):
    beta = math.radians(beta)      #仰角
    R_earth = 6371                  #地球半径
    eq_R_earth = (4.0 / 3.0) * R_earth  #km 等效地球半径
    h_0 = 1.079                         #km雷达海拔高度
    a = d / eq_R_earth                  #地心角
    # alpha = (4.0 / 3.0) * a             #等效地心角
    alpha = a                          #好像这样是对的
    # r = 1
    r = 0.0625                    # 库长 62.5m
    G = ((eq_R_earth + h_0) * np.sin(alpha)) / (r * np.cos(alpha + beta))

    return G

def calculate_distance_shifts(distances, elvs):
    """
    计算不同仰角层的距离库偏移。
    参数:
        distances (list): 距离范围（单位：km）。
        elvs (list): 仰角列表（单位：度）。
    返回:
        dict: 每个仰角对应的距离库偏移。
    """
    shifts = {}
    for elv in elvs:
        shifts[elv] = [int(cal_radar_gauge(d, elv)) for d in distances]
    return shifts

def calculate_relative_shifts(distances, elvs, reference_elv=0.5):
    """
    计算相对库偏移。
    参数:
        distances (list): 距离范围（单位：km）。
        elvs (list): 仰角列表（单位：度）。
        reference_elv (float): 参考仰角（默认最低仰角）。
    返回:
        dict: 每个仰角相对参考仰角的库偏移。
    """
    reference_shifts = [int(cal_radar_gauge(d, reference_elv)) for d in distances]
    shifts = {}
    for elv in elvs:
        current_shifts = [int(cal_radar_gauge(d, elv)) for d in distances]
        shifts[elv] = [current - ref for current, ref in zip(current_shifts, reference_shifts)]
    return shifts

#分仰角层输出一个时次的雷达数据
def process_radar_data(dataarray, layer_n):  # 输出第几层的数据，从第一层开始 第一层:0.5°仰角
    layers = []
    current_layer = []
    prev_angle = -np.inf

    for idx, row in dataarray.iterrows():
        angle = row[0]
        # 提取每行的第一个值，即为方位角度数
        if angle < prev_angle:  # 若切换则存储该层到layers列表中
            layers.append(pd.DataFrame(current_layer))  # 判断是否切换层
            current_layer = []  # 若切换则存储该层到layers列表中

        current_layer.append(row)  # 没切换层则添加到当前层列表中
        prev_angle = angle
        # 存储到过去方位角数据，循环后与下一个方位角比较用来判断是否切换层

    if current_layer:
        layers.append(pd.DataFrame(current_layer))

    return layers, layers[layer_n - 1]

# 把前一个函数输出的数据对应多层输出同一径向区间数据
# def mix_layers(layers, target_layers):
#     angle_ranges = [
#         (12000.0, 19000.0),  # 第一段
#         (22500.0, 23500.0),  # 第二段
#         (25500.0, 26500.0)  # 第三段
#     ]
#
#     # 初始化结果容器
#     segmented_data = [[] for _ in range(len(angle_ranges))]
#
#     # 遍历目标层
#     for layer_idx in target_layers:
#         layer = layers[layer_idx]
#
#         # 遍历每个角度范围
#         for range_idx, (start_angle, end_angle) in enumerate(angle_ranges):
#             # 筛选当前角度范围的数据
#             selected_rows = layer[(layer.iloc[:, 0] >= start_angle) & (layer.iloc[:, 0] <= end_angle)]
#             if not selected_rows.empty:
#                 segmented_data[range_idx].append(selected_rows.values)
#
#     # 将每段数据转换为三维数组
#     segmented_arrays = [np.array(data) for data in segmented_data]
#
#     return segmented_arrays

def align_radar_data(dataarray, angle_ranges):
    """
    对齐不同仰角的雷达数据，使得它们在每个范围内的径向数相同。

    Args:
        dataarray (pd.DataFrame): 包含雷达数据的 DataFrame，每行是一条径向数据，第一列为仰角。
        angle_ranges (list of tuple): 方位角范围列表，每个范围是一个元组 (start, end)。

    Returns:
        list of list: 对齐后的所有仰角数据，每个仰角是一个 DataFrame 列表。
    """

    # 提取所有仰角数据
    layers, _ = process_radar_data(dataarray, 1)  # 提取所有层
    all_layers_data = [process_radar_data(dataarray, i + 1)[1] for i in range(len(layers))]

    # 计算每个范围的最小径向数
    selected_rows_shapes = []
    for angle_range in angle_ranges:
        min_rows = np.inf
        for layer_data in all_layers_data:
            # 统计该范围内的径向数
            rows_in_range = layer_data[
                (layer_data.iloc[:, 0] >= angle_range[0]) & (layer_data.iloc[:, 0] <= angle_range[1])
            ]
            min_rows = min(min_rows, len(rows_in_range))
        selected_rows_shapes.append(min_rows)

    # 对齐数据
    aligned_data = []
    for layer_data in all_layers_data:
        aligned_layer = []
        for range_idx, angle_range in enumerate(angle_ranges):
            rows_in_range = layer_data[
                (layer_data.iloc[:, 0] >= angle_range[0]) & (layer_data.iloc[:, 0] <= angle_range[1])
            ]
            if len(rows_in_range) > selected_rows_shapes[range_idx]:
                # 根据第一列的方位角值，选择最接近的数据
                closest_rows = rows_in_range.iloc[
                    np.argsort(
                        np.abs(
                            rows_in_range.iloc[:, 0].values
                            - np.median(rows_in_range.iloc[:, 0].values)
                        )
                    )[:selected_rows_shapes[range_idx]]
                ]
                aligned_layer.append(closest_rows.sort_index())
            else:
                aligned_layer.append(rows_in_range)
        aligned_data.append(aligned_layer)

    return aligned_data

def extract_to_ndarray(aligned_data, selected_layers, selected_range_idx):
    """
    从对齐后的数据中提取指定层和方位范围的数据，转换为三维 ndarray。

    Args:
        aligned_data (list of list): 对齐后的数据，每个仰角包含多个方位范围的数据。
        selected_layers (list of int): 要提取的层索引列表（从 0 开始）。
        selected_range_idx (int): 要提取的方位范围索引（从 0 开始）。

    Returns:
        np.ndarray: 提取后的三维数组，形状为 (len(selected_layers), rows, columns)。
    """
    extracted_data = []

    for layer_idx in selected_layers:
        # 获取指定层的数据
        layer_data = aligned_data[layer_idx][selected_range_idx]
        # 将 DataFrame 转换为 ndarray 并添加到结果中
        extracted_data.append(layer_data.values)

    # 将列表转换为三维 ndarray
    return np.array(extracted_data)

# 求中间仰角的按权重计算平均值
# def interpolate_layer(data, lower_idx, upper_idx, new_idx, new_angle, lower_angle, upper_angle):
#     lower_layer = data[lower_idx]
#     upper_layer = data[upper_idx]
#
#     weight = (new_angle - lower_angle) / (upper_angle - lower_angle)
#
#     interpolate_layer = lower_layer * (1 - weight) + upper_layer * weight
#
#     data[new_idx] = interpolate_layer
#
#     data = np.delete(data, upper_idx, axis = 0)
#
#     return data

# 3*3数据块筛选提取

def interpolate_layer(data, lower_idx, upper_idx, new_idx, new_angle, lower_angle, upper_angle):
    """
    线性插值函数，考虑无效值(-32768)的处理

    Args:
        data: 输入的数据数组
        lower_idx: 较低仰角层的索引
        upper_idx: 较高仰角层的索引
        new_idx: 新插值层的索引
        new_angle: 新插值层的仰角
        lower_angle: 较低层的仰角
        upper_angle: 较高层的仰角

    Returns:
        处理后的数据数组
    """
    lower_layer = data[lower_idx]
    upper_layer = data[upper_idx]

    # 创建掩码标记无效值位置
    invalid_mask = (lower_layer == -32768) | (upper_layer == -32768)

    # 计算权重
    weight = (new_angle - lower_angle) / (upper_angle - lower_angle)

    # 创建插值结果数组
    interpolated_layer = np.zeros_like(lower_layer)

    # 对有效值进行插值
    valid_mask = ~invalid_mask
    interpolated_layer[valid_mask] = (lower_layer[valid_mask] * (1 - weight) +
                                      upper_layer[valid_mask] * weight)

    # 将无效值位置设为-32768
    interpolated_layer[invalid_mask] = -32768

    # 更新数据
    data[new_idx] = interpolated_layer

    # 删除较高仰角层
    data = np.delete(data, upper_idx, axis=0)

    return data
def extract_3x3_matrix(layer_data, i, j, grid_size=3):
    """
    提取指定点 (i, j) 周围的 3x3 矩阵。
    """
    half_size = grid_size // 2
    if (i - half_size < 0 or i + half_size >= layer_data.shape[0] or
        j - half_size < 0 or j + half_size >= layer_data.shape[1]):
        return None  # 跳过边界点
    return layer_data[i - half_size:i + half_size + 1, j - half_size:j + half_size + 1]

# def create_dataset(layer_arrays, distance_shifts, max_bins=3200, grid_size=3):
#     """
#     构建数据集。
#     参数:
#         layer_arrays (list): 仰角层数据。
#         distance_shifts (dict): 各仰角层的距离库偏移。
#         max_bins (int): 最大距离库数。
#         grid_size (int): 矩阵大小。
#     返回:
#         tuple: (特征数据, 标签数据)。
#     """
#     feature_data = []
#     label_data = []
#
#     # 标签层为第2层（0.89°）
#     target_layer = layer_arrays[1]
#     for i in range(grid_size // 2, target_layer.shape[0] - grid_size // 2):
#         for j in range(grid_size // 2, min(target_layer.shape[1] - grid_size // 2, max_bins)):
#             # 提取标签值
#             label = target_layer[i, j]
#
#             # 跳过无效点
#             if label == -32768:
#                 continue
#
#             # 提取自变量（0.5°、1.45°、2.4°层的 3x3 矩阵）
#             feature_vector = []
#             for layer_idx, elv in zip([0, 2, 3], [0.5, 1.45, 2.4]):
#                 shifted_j = j + distance_shifts[elv][j]  # 根据距离库差调整列索引
#                 matrix = extract_3x3_matrix(layer_arrays[layer_idx], i, shifted_j, grid_size)
#                 if matrix is not None:
#                     feature_vector.append(matrix)
#                 else:
#                     break  # 如果某层的矩阵无效，则跳过该点
#             else:
#                 # 保存特征和标签
#                 feature_data.append(feature_vector)
#                 label_data.append(label)
#
#     return np.array(feature_data), np.array(label_data)

# def create_dataset(layer_arrays, distance_shifts, max_bins=3200, grid_size=3):
#     """
#     构建数据集。
#     参数:
#         layer_arrays (list): 仰角层数据。
#         distance_shifts (dict): 各仰角层的距离库偏移（相对）。
#         max_bins (int): 最大距离库数。
#         grid_size (int): 矩阵大小。
#     返回:
#         tuple: (特征数据, 标签数据)。
#     """
#     feature_data = []
#     label_data = []
#
#     # 标签层为第2层（0.89°）
#     # target_layer = layer_arrays[1]
#     target_layer = layer_arrays[1]
#     for i in range(grid_size // 2, target_layer.shape[0] - grid_size // 2):
#         for j in range(grid_size // 2, min(target_layer.shape[1] - grid_size // 2, max_bins)):
#             # 提取标签值
#             label = target_layer[i, j]
#
#             # 跳过无效点
#             if label == -32768:
#                 continue
#
#             # 提取自变量（0.5°、1.45°、2.4°层的 3x3 矩阵）
#             feature_vector = []
#             # for layer_idx, elv in zip([0, 2, 3], [0.5, 1.45, 2.4]):
#             for layer_idx, elv in zip([0, 2, 3], [2.4, 3.5, 4.3]):
#                 # 使用相对库差调整列索引
#                 shifted_j = j + distance_shifts[elv][j]
#                 if shifted_j < 0 or shifted_j >= target_layer.shape[1]:
#                     break  # 超出范围
#                 matrix = extract_3x3_matrix(layer_arrays[layer_idx], i, shifted_j, grid_size)
#                 if matrix is not None:
#                     feature_vector.append(matrix)
#                 else:
#                     break  # 如果某层的矩阵无效，则跳过该点
#                 # 检查3x3矩阵中是否包含无效值
#                 # if matrix is None or np.any(matrix == -32768):
#                 #     break
#                 # feature_vector.append(matrix)
#             else:
#                 # 保存特征和标签
#                 feature_data.append(feature_vector)
#                 label_data.append(label)
#
#     return np.array(feature_data), np.array(label_data)
def create_dataset(layer_arrays, distance_shifts, layer_indices, elevations, target_layer_idx, max_bins=3200,
                   grid_size=3):
    """
    构建数据集。

    Args:
        layer_arrays: 雷达数据数组
        distance_shifts: 距离库偏移字典
        layer_indices: 要使用的层索引列表，如 [0, 2, 3]
        elevations: 对应的仰角列表，如 [0.5, 1.45, 2.4]
        target_layer_idx: 目标层索引
        max_bins: 最大距离库数
        grid_size: 特征矩阵大小

    Returns:
        tuple: (特征数据, 标签数据)
    """
    feature_data = []
    label_data = []

    # 统计计数器
    total_points = 0
    invalid_labels = 0
    invalid_matrices = 0
    valid_samples = 0
    out_of_range = 0

    target_layer = layer_arrays[target_layer_idx]

    # 检查输入参数
    if len(layer_indices) != len(elevations):
        raise ValueError("层索引列表和仰角列表长度必须相同")

    # 遍历数据点
    for i in range(grid_size // 2, target_layer.shape[0] - grid_size // 2):
        for j in range(grid_size // 2, min(target_layer.shape[1] - grid_size // 2, max_bins)):
            total_points += 1
            label = target_layer[i, j]

            # 检查标签
            if label == -32768:
                invalid_labels += 1
                continue

            # 提取特征
            feature_vector = []
            valid_feature = True

            for layer_idx, elv in zip(layer_indices, elevations):
                shifted_j = j + distance_shifts[elv][j]

                # 检查索引范围
                if shifted_j < 0 or shifted_j >= target_layer.shape[1]:
                    out_of_range += 1
                    valid_feature = False
                    break

                matrix = extract_3x3_matrix(layer_arrays[layer_idx], i, shifted_j, grid_size)

                # 检查矩阵有效性
                if matrix is None or np.any(matrix == -32768):
                    invalid_matrices += 1
                    valid_feature = False
                    break

                feature_vector.append(matrix)

            if valid_feature:
                feature_data.append(feature_vector)
                label_data.append(label)
                valid_samples += 1

    return np.array(feature_data), np.array(label_data)


def load_and_process_data(root_path, filename_ZH, filename_CC):
    file_path_ZH = root_path + filename_ZH
    file_path_CC = root_path + filename_CC

    # 加载数据
    data_ZH = pd.read_csv(file_path_ZH, header=None)
    data_CC = pd.read_csv(file_path_CC, header=None)

    # 利用 cc 数据质控
    filtered_zh = data_ZH.where(data_CC > 0.8, -32768)

    # 分成按层排列的数据
    # layers, layer_1 = process_radar_data(filtered_zh, 1)

    # 归整成三维数组并且计算中间层的线性插值
    angle_ranges = [
        (12000.0, 19000.0),  # 第一段
        (22500.0, 23500.0),  # 第二段
        (25500.0, 26500.0)  # 第三段
    ]

    aligned_data = align_radar_data(filtered_zh, angle_ranges)
    # selected_layers = [0, 1, 2, 3, 4]
    # selected_layers = [4, 5, 6, 7, 8]
    selected_layers = [7, 8, 9, 10, 11, 12, 13, 14] #4.0,4.5,5.0,6.0,7.0,8.0,9.0,9.9
    selected_range_idx = 0  # 选择第 2 个方位范围（索引从 0 开始）

    array = extract_to_ndarray(aligned_data, selected_layers, selected_range_idx)
    array = array[:, :, 1:]
    # array = np.round(interpolate_layer(array, 2, 3, 2, 1.45, 1.29, 1.79))
    # array = np.round(interpolate_layer(array, 3, 4, 3, 4.3, 4.0, 4.5))
    array = np.round(interpolate_layer(array, 0, 1, 0, 4.3, 4.0, 4.5))
    print(array.shape)
    # 提取标签数据和特征数据
    # distances = np.linspace(0, 200, 3200)
    # elvs = [0.5, 0.89, 1.45, 2.4]  # 仰角
    # distance_shifts = calculate_distance_shifts(distances, elvs)
    all_distances = np.linspace(0, 200, 3200)  # 全部距离范围
    selected_distances = all_distances[80:1600]  # 选择第 80 到第 1600 个库
    # elvs = [0.5, 0.89, 1.45, 2.4]
    # elvs = [2.4, 3.0, 3.5, 4.3]
    elvs = [4.3, 5.0, 6.0, 7.0, 8.0, 9.0, 9.9]
    distance_shifts = calculate_relative_shifts(selected_distances, elvs)

    # 使用 create_dataset 提取特征和标签
    feature_data, label_data = create_dataset(
        array,
        distance_shifts,
        layer_indices=[0,2,6],
        # elevations=[0.5,1.45,2.4],
        # layer_indices=[4, 6, 7],
        # elevations=[2.4, 3.5, 4.3],
        elevations=[4.3, 6.0, 9.9],
        target_layer_idx=5,
        max_bins=1520
    )

    return feature_data, label_data

def process_interpolation_data(root_path, filename_ZH, filename_CC):
    file_path_ZH = root_path + filename_ZH
    file_path_CC = root_path + filename_CC

    # 加载数据
    data_ZH = pd.read_csv(file_path_ZH)
    data_CC = pd.read_csv(file_path_CC)

    # 利用 cc 数据质控
    filtered_zh = data_ZH.where(data_CC > 0.8, -32768)

    # 分成按层排列的数据
    # layers, layer_1 = process_radar_data(filtered_zh, 1)

    # 归整成三维数组并且计算中间层的线性插值
    angle_ranges = [
        (12000.0, 19000.0),  # 第一段
        (22500.0, 23500.0),  # 第二段
        (25500.0, 26500.0)  # 第三段
    ]

    aligned_data = align_radar_data(filtered_zh, angle_ranges)
    selected_layers = [0, 1, 2, 3, 4]  # 选择第 1、6、11 层（索引从 0 开始）
    selected_range_idx = 0  # 选择第 2 个方位范围（索引从 0 开始）

    array = extract_to_ndarray(aligned_data, selected_layers, selected_range_idx)
    array = array[:, :, 1:]
    array = np.round(interpolate_layer(array, 2, 3, 2, 1.45, 1.29, 1.79))

    return array